﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using OpenCL.Net;
using System.IO;
using SharpRBM.Core.Enumerations;

namespace SharpRBM.Core.LearningModules
{
    public class OpenCLLearningModule : ILearningModule
    {
        private const int NUMBER_OF_RANDOMS = 7919;
        public const string ProgramFileName = @"LearningModules\TrainRBM.c";
        public const string PrngFileName = @"LearningModules\PsuedoRandom.c";
        public const string DevProgramFileName = @"..\..\..\SharpRBM.Core\LearningModules\TrainRBM.c";
        public const string DevPrngFileName = @"..\..\..\SharpRBM.Core\LearningModules\PsuedoRandom.c";
        private Cl.Context _context;
        private Cl.Device _device;
        private float[] _detectorError;
        private float[] _reconstructionError;
        private float[] _model;
        private float[] _reconstructed;
        private float[] _reconstructedModel;
        private float[] _twoPhaseBuffer;
        private int _testCount;
        private List<Layer> _inputs;
        private InterLayerWeights _weights;
        private Cl.Kernel _clearDetectorErrors;
        private Cl.Kernel _updateWeights;

        private Cl.Kernel _activateLowerToUpper;
        private Cl.Kernel _activateLowerToUpperBinary;
        private Cl.Kernel _activateUpperToLower;
        private Cl.Kernel _accumulateErrors;
        private Cl.Kernel _sumErrors;
        private Cl.Kernel _sumErrorsDetector;
        private Cl.Kernel _sumErrorsRecon;

        private Cl.CommandQueue _commandQueue;
        private Cl.Mem _weightsMem;
        private Cl.Mem _detectorErrorMem;
        private Cl.Mem _epsilonMem;
        private Cl.Mem _modelMem;
        private Cl.Mem _reconstructedMem;
        private Cl.Mem _reconstructionErrorMem;
        private Cl.Mem _reconstructedModelMem;
        private Cl.Mem _errorSumsMem;
        private Cl.Mem _twoPhaseMem;

        private List<Cl.Mem> _testCaseMems;
        private int _inPtrSize = 0;
        private float[] _epsilon = new float[1];
        private List<Cl.Mem> _mems = new List<Cl.Mem>();
        private List<Cl.Kernel> _kernels = new List<Cl.Kernel>();

        public OpenCLLearningModule()
        {
            unsafe
            {
                _inPtrSize = sizeof(IntPtr);
            }
        }

        public string GetDeviceInfo()
        {
            Cl.ErrorCode error;
            IntPtr size = Cl.GetDeviceInfo(_device, Cl.DeviceInfo.GlobalMemSize, out error).CastTo<IntPtr>();
            return size.ToInt32().ToString();
        }

        public void Prepare(InterLayerWeights weights, List<Layer> inputs)
        {
            _inputs = inputs;
            GetDevice();
            GetContext();
            _weights = weights;
            _model = new float[weights.UpperLayerSize];
            _reconstructed = new float[weights.LowerLayerSize];
            _reconstructionError = new float[_weights.Weights.Length];
            _reconstructedModel = new float[weights.LowerLayerSize];
            _twoPhaseBuffer = new float[weights.LowerLayerSize > weights.UpperLayerSize ? weights.LowerLayerSize : weights.UpperLayerSize];
            _testCount = inputs.Count;
            _detectorError = new float[weights.UpperLayerSize * weights.LowerLayerSize];
            LoadProgram();
            _weightsMem.WriteFloatArray(_commandQueue, _weights.Weights, BlockMode.Blocking);
        }

        public void ActivateLowerToUpper(Layer lower, Layer upper, ActivationMethod method)
        {
            //_weights.ActivateLowerToUpper(lower, upper, method);
            //return;
            Cl.Mem testCaseMem = _context.AllocateFloatArray(lower.Values, Cl.MemFlags.CopyHostPtr | Cl.MemFlags.ReadOnly);
            try
            {
                testCaseMem.WriteFloatArray(_commandQueue, lower.Values, BlockMode.Blocking);

                ( method == ActivationMethod.Binary ? _activateLowerToUpperBinary : _activateLowerToUpper )              
                    .SetKernelArgFloat(0, testCaseMem)
                    .SetKernelArgFloat(1, _reconstructedModelMem)
                    .SetKernelArgFloat(2, _weightsMem)
                    .SetKernelArg( 3, Helper.Random.Next() )
                    .EnqueueNDRangeKernel( _commandQueue, _weights.UpperLayerSize );
                _reconstructedModelMem.ReadFloatArray(_commandQueue, upper.Values, BlockMode.Blocking);
            }
            finally
            {
                Cl.ReleaseMemObject(testCaseMem);
            }
        }

        public void ActivateUpperToLower(Layer lower, Layer upper)
        {
            //_weights.ActivateUpperToLower(lower, upper);
            //return;
            _modelMem.WriteFloatArray(_commandQueue, upper.Values, BlockMode.Blocking);

            _activateUpperToLower
                .SetKernelArgFloat(0, _reconstructedMem)
                .SetKernelArgFloat(1, _modelMem)
                .SetKernelArgFloat(2, _weightsMem)
                .EnqueueNDRangeKernel(_commandQueue, _weights.LowerLayerSize);

            _reconstructedMem.ReadFloatArray(_commandQueue, lower.Values, BlockMode.Blocking);
        }

        public TrainingError Train()
        {
            TrainingError trainingError = new TrainingError();

            List<Cl.Mem> randomizedTestMemCases = _testCaseMems.ToList();
            randomizedTestMemCases.Shuffle();

            // Update weights
            _epsilon[0] = _weights.GetAdjustedAndScaledTrainingRate(_testCount);
            _epsilonMem.WriteFloatArray(_commandQueue, _epsilon, BlockMode.Blocking);

            Helper.BatchForEach(randomizedTestMemCases, 10, testCases =>
                {
                    // Make sure we're done with everything that's been requested before
                    _commandQueue.Finish();

                    _clearDetectorErrors
                        .SetKernelArgFloat(0, _detectorErrorMem)
                        .SetKernelArgFloat(1, _reconstructionErrorMem)
                        .EnqueueNDRangeKernel(_commandQueue, _weights.LowerLayerSize, _weights.UpperLayerSize);

                    foreach (Cl.Mem testCase in testCases)
                    {
                        _activateLowerToUpperBinary
                            .SetKernelArgFloat(0, testCase)
                            .SetKernelArgFloat(1, _modelMem)
                            .SetKernelArgFloat(2, _weightsMem)
                            .SetKernelArg( 3, Helper.Random.Next() )
                            .EnqueueNDRangeKernel( _commandQueue, _weights.UpperLayerSize );

                        _activateUpperToLower
                            .SetKernelArgFloat(0, _reconstructedMem)
                            .SetKernelArgFloat(1, _modelMem)
                            .SetKernelArgFloat(2, _weightsMem)
                            .EnqueueNDRangeKernel(_commandQueue, _weights.LowerLayerSize);

                        _activateLowerToUpper
                            .SetKernelArgFloat(0, _reconstructedMem)
                            .SetKernelArgFloat(1, _reconstructedModelMem)
                            .SetKernelArgFloat(2, _weightsMem)
                            .SetKernelArg( 3, Helper.Random.Next() )
                            .EnqueueNDRangeKernel(_commandQueue, _weights.UpperLayerSize);

                        _accumulateErrors
                            .SetKernelArgFloat(0, testCase)
                            .SetKernelArgFloat(1, _detectorErrorMem)
                            .SetKernelArgFloat(2, _modelMem)
                            .SetKernelArgFloat(3, _reconstructedMem)
                            .SetKernelArgFloat(4, _reconstructedModelMem)
                            .SetKernelArgFloat(5, _reconstructionErrorMem)
                            .EnqueueNDRangeKernel(_commandQueue, _weights.LowerLayerSize, _weights.UpperLayerSize);
                    }

                    _updateWeights
                        .SetKernelArgFloat(0, _weightsMem)
                        .SetKernelArgFloat(1, _detectorErrorMem)
                        .SetKernelArgFloat(2, _epsilonMem)
                        .EnqueueNDRangeKernel(_commandQueue, _weights.LowerLayerSize, _weights.UpperLayerSize);

                    float[] errorSums = new float[2];

                    _sumErrorsDetector
                        .SetKernelArgFloat( 0, _detectorErrorMem )
                        .SetKernelArgFloat( 1, _reconstructionErrorMem )
                        .SetKernelArgFloat( 2, _errorSumsMem )
                        .SetKernelArgFloat( 3, _twoPhaseMem )
                        .EnqueueNDRangeKernel( _commandQueue, _weights.LowerLayerSize );

                    _sumErrorsRecon
                        .SetKernelArgFloat( 0, _detectorErrorMem )
                        .SetKernelArgFloat( 1, _reconstructionErrorMem )
                        .SetKernelArgFloat( 2, _errorSumsMem )
                        .SetKernelArgFloat( 3, _twoPhaseMem )
                        .EnqueueNDRangeKernel( _commandQueue, _weights.LowerLayerSize );


                    _errorSumsMem.ReadFloatArray(_commandQueue, errorSums, BlockMode.Blocking);

                    trainingError.FeatureDetectorError += errorSums[0];
                    trainingError.ReconstructionError += errorSums[1];
                });

            // copy results from device back to host
            _weightsMem.ReadFloatArray(_commandQueue, _weights.Weights, BlockMode.Blocking);

            // Make sure we're done with everything that's been requested before
            _commandQueue.Finish();

            // TODO: Wait for the event that should be returned by _errorSumsMem.ReadFloatArray...                
            return trainingError;
        }

        private IntPtr[] GetLowerUpperDimensions()
        {
            return new IntPtr[] { new IntPtr(_weights.LowerLayerSize), new IntPtr(_weights.UpperLayerSize) };
        }

        public void Dispose()
        {
            _mems.ForEach(mem => Cl.ReleaseMemObject(mem));
            _kernels.ForEach(kernel => Cl.ReleaseKernel(kernel));
            Cl.ReleaseCommandQueue(_commandQueue);
        }

        private void GetDevice()
        {
            Cl.ErrorCode error;

            var platforms = Cl.GetPlatformIDs(out error);
            if (error != Cl.ErrorCode.Success) {
                throw new InvalidOperationException(string.Format("Unable to retrieve an OpenCL Device, error was: {0}!", error));
            }


            Cl.Device? device = null;

            //CHANGE THIS TO CHANGE THE RELATIVE PRIORITY OF CPU / GPU
            Cl.DeviceType[] devicePriority = { Cl.DeviceType.Cpu, Cl.DeviceType.Gpu, Cl.DeviceType.Accelerator, Cl.DeviceType.All };

            foreach (var deviceType in devicePriority)
            {
                //look for GPUs first
                foreach (var platform in platforms)
                {
                    var deviceIds = Cl.GetDeviceIDs(platform, deviceType, out error);
                    if (deviceIds.Any())
                    {
                        device = deviceIds.First();
                        break;
                    }
                }
            }

            _device = device.Value;

            if (error != Cl.ErrorCode.Success)
            {
                throw new InvalidOperationException(string.Format("Unable to retrieve an OpenCL Device, error was: {0}!", error));
            }
        }

        private void GetContext()
        {
            Cl.ErrorCode error;

            _context = Cl.CreateContext(null, 1, new[] { _device }, null, IntPtr.Zero, out error);

            if (error != Cl.ErrorCode.Success)
            {
                throw new InvalidOperationException(string.Format("Unable to retrieve an OpenCL Context, error was: {0}!", error));
            }
        }

        private void LoadProgram()
        {
            string source = LoadSourceCode();

            Cl.ErrorCode error;

            using (Cl.Program program = Cl.CreateProgramWithSource(_context, 1, new[] { source }, null, out error))
            {
                if (error != Cl.ErrorCode.Success)
                {
                    throw new InvalidProgramException(string.Format("Unable to run Cl.CreateProgramWithSource for program: {0}!", error));
                }

                error = Cl.BuildProgram(program, 1, new[] { _device }, string.Empty, null, IntPtr.Zero);
                if (error != Cl.ErrorCode.Success)
                {
                    char[] buffer = new char[90000];
                    Cl.InfoBuffer infoBuffer = new Cl.InfoBuffer(new IntPtr(90000));
                    IntPtr retSize;
                    Cl.GetProgramBuildInfo(program, _device, Cl.ProgramBuildInfo.Log, new IntPtr(90000), infoBuffer, out retSize);

                    throw new InvalidProgramException("Error building program:\n" + infoBuffer.ToString());
                }

                error = Cl.BuildProgram(program, 1, new[] { _device }, string.Empty, null, IntPtr.Zero);
                if (error != Cl.ErrorCode.Success)
                {
                    throw new InvalidProgramException(string.Format("Unable to run Cl.BuildProgram for program: {0}!", error));
                }

                Cl.BuildStatus buildStatus = Cl.GetProgramBuildInfo(program, _device, Cl.ProgramBuildInfo.Status, out error).CastTo<Cl.BuildStatus>();
                if (buildStatus != Cl.BuildStatus.Success)
                {
                    throw new InvalidProgramException(string.Format("GetProgramBuildInfo returned {0} for program!", buildStatus));
                }

                _commandQueue = Cl.CreateCommandQueue(_context, _device, (Cl.CommandQueueProperties)0, out error);
                if (error != Cl.ErrorCode.Success)
                {
                    throw new InvalidProgramException(string.Format("Unable to CreateCommandQueue: {0}!", error));
                }

                GenerateStorage();

                _accumulateErrors = CreateKernel(program, "AccumulateErrors");
                _clearDetectorErrors = CreateKernel(program, "ClearDetectorErrors");

                // cj: Moved constants out of CL code to make it more concise and expressive and easier to touch
                // Since a recompile was required, I just used comments here with different (albiet poorly named) kernels
                // At least these can be switch at runtime in the future with litte effort
                _activateLowerToUpperBinary = CreateKernel( program, "ActivateLowerToUpperBinary_Binary" );
                //_activateLowerToUpperBinary = CreateKernel( program, "ActivateLowerToUpperBinary_Linear" );
                //_activateLowerToUpper = CreateKernel( program, "ActivateLowerToUpper_Binary" );
                _activateLowerToUpper = CreateKernel( program, "ActivateLowerToUpper_Linear" );

                //_activateLowerToUpperBinary = CreateKernel( program, "ActivateLowerToUpperBinary" );
                //_activateLowerToUpper = CreateKernel( program, "ActivateLowerToUpper" );

                _activateUpperToLower = CreateKernel( program, "ActivateUpperToLower" );
                _updateWeights = CreateKernel(program, "UpdateWeights");
                _sumErrorsDetector = CreateKernel( program, "SumErrorsDetector" );
                _sumErrorsRecon = CreateKernel( program, "SumErrorsRecon" );
            }
        }


        private Cl.Kernel CreateKernel(Cl.Program program, string kernelName)
        {
            Cl.Kernel kernel = program.CreateKernel(kernelName);
            _kernels.Add(kernel);
            return kernel;
        }

        private void GenerateStorage()
        {
            IntPtr lowerCount = new IntPtr(_weights.LowerLayerSize);
            IntPtr upperCount = new IntPtr(_weights.UpperLayerSize);
            IntPtr testCount = new IntPtr(_testCount);

            _weightsMem = AllocateFloatArray(_weights.Weights, Cl.MemFlags.CopyHostPtr | Cl.MemFlags.ReadWrite);
            _detectorErrorMem = AllocateFloatArray(_detectorError, Cl.MemFlags.CopyHostPtr | Cl.MemFlags.ReadWrite);
            _epsilonMem = AllocateFloatArray(_epsilon, Cl.MemFlags.CopyHostPtr | Cl.MemFlags.ReadOnly);
            _modelMem = AllocateFloatArray(_model, Cl.MemFlags.CopyHostPtr | Cl.MemFlags.ReadWrite);
            _reconstructedMem = AllocateFloatArray(_reconstructed, Cl.MemFlags.CopyHostPtr | Cl.MemFlags.ReadWrite);
            _reconstructedModelMem = AllocateFloatArray(_reconstructedModel, Cl.MemFlags.CopyHostPtr | Cl.MemFlags.ReadWrite);
            _reconstructionErrorMem = AllocateFloatArray(_reconstructionError, Cl.MemFlags.CopyHostPtr | Cl.MemFlags.ReadWrite);
            _errorSumsMem = AllocateFloatArray(new float[2], Cl.MemFlags.CopyHostPtr | Cl.MemFlags.ReadWrite);
            _twoPhaseMem = AllocateFloatArray( _twoPhaseBuffer, Cl.MemFlags.CopyHostPtr | Cl.MemFlags.ReadWrite );

            _testCaseMems = new List<Cl.Mem>();
            for (int testCase = 0; testCase < _testCount; testCase++)
            {
                Cl.Mem testCaseMem = AllocateFloatArray(_inputs[testCase].Values, Cl.MemFlags.CopyHostPtr | Cl.MemFlags.ReadOnly);
                testCaseMem.WriteFloatArray(_commandQueue, _inputs[testCase].Values, BlockMode.NonBlocking);
                _testCaseMems.Add(testCaseMem);
            }
        }

        private Cl.Mem AllocateFloatArray(float[] array, Cl.MemFlags flags)
        {
            Cl.Mem result = _context.AllocateFloatArray(array, flags);
            _mems.Add(result);
            return result;
        }

        private Cl.Mem AllocateIntArray(int[] array, Cl.MemFlags flags)
        {
            Cl.Mem result = _context.AllocateIntArray(array, flags);
            _mems.Add(result);
            return result;
        }

        private string LoadSourceCode()
        {
            // If we don't do a build, the source code isn't copied, therefore, for developping purposes, we try to load it from the dev folder first

            string source;

            try
            {
                source = File.ReadAllText( DevPrngFileName );
                source += File.ReadAllText( DevProgramFileName );
            }
            catch
            {
                source = File.ReadAllText( PrngFileName );
                source += File.ReadAllText( ProgramFileName );
            }

            source =
                source
                    .Replace("<UpperCount>", _weights.UpperLayerSize.ToString())
                    .Replace("<LowerCount>", _weights.LowerLayerSize.ToString())
                    .Replace("<RandomCount>", NUMBER_OF_RANDOMS.ToString());

            return source;
        }
    }
}
